"""
下载模块：处理视频下载相关的功能
包括视频解析、下载、合并等操作
"""

import os
import uuid
import asyncio
import time
import aiohttp
import re
from pathlib import Path
from datetime import datetime
from typing import Dict, Tuple, Optional, List, Union

from bbdown.core.parser import Parser
from bbdown.core.config import Config
from bbdown.core.logger import Logger
from bbdown.utils.downloader import Downloader
from bbdown.utils.merger import Merger
from bbdown.utils.danmaku import DanmakuUtil
from bbdown.utils.subtitle import SubtitleUtil

from .constants import DOWNLOAD_DIR, TASKS
from .db import insert_task, update_task_status, insert_result

def update_progress(task_id: str, part: str, downloaded: int, total: int):
    """
    更新下载进度
    
    参数：
        task_id (str): 任务ID
        part (str): 下载部分（video/audio）
        downloaded (int): 已下载字节数
        total (int): 总字节数
    """
    if task_id in TASKS:
        if part == "video":
            TASKS[task_id]["video_progress"] = (downloaded, total)
            TASKS[task_id]["video_speed"] = downloaded / (time.time() - TASKS[task_id].get("start_time", time.time()))
        elif part == "audio":
            TASKS[task_id]["audio_progress"] = (downloaded, total)
            TASKS[task_id]["audio_speed"] = downloaded / (time.time() - TASKS[task_id].get("start_time", time.time()))
        
        # 计算总体进度
        video_progress = TASKS[task_id].get("video_progress", (0, 1))
        audio_progress = TASKS[task_id].get("audio_progress", (0, 1))
        
        total_progress = (
            (video_progress[0] + audio_progress[0]) /
            (video_progress[1] + audio_progress[1]) * 100
        )
        TASKS[task_id]["progress"] = int(total_progress)
        
        # 计算总速度
        total_speed = TASKS[task_id].get("video_speed", 0) + TASKS[task_id].get("audio_speed", 0)
        TASKS[task_id]["total_speed"] = total_speed

def init_download_task(task_id: str) -> None:
    """
    初始化下载任务状态
    
    参数：
        task_id (str): 任务ID
    """
    TASKS[task_id] = {
        "status": "running",
        "progress": 0,
        "current_step": "正在解析视频信息",
        "video_progress": (0, 1),
        "audio_progress": (0, 1),
        "video_speed": 0,
        "audio_speed": 0,
        "total_speed": 0,
        "start_time": time.time()
    }
    update_task_status(task_id, "running")

async def parse_bvid(url: str) -> str:
    """
    从URL中解析bvid
    
    参数：
        url (str): 视频URL
        
    返回：
        str: 视频ID
    """
    if "b23.tv" in url:
        # 处理短链接
        async with aiohttp.ClientSession() as session:
            async with session.get(url, allow_redirects=True) as response:
                url = str(response.url)
    
    if "bilibili.com/video/" in url:
        bvid = url.split("bilibili.com/video/")[1].split("?")[0].split("/")[0]
    elif "b23.tv/" in url:
        bvid = url.split("b23.tv/")[1].split("?")[0].split("/")[0]
    else:
        raise ValueError("无法从URL中解析出视频ID")
    
    return bvid

def find_part(video_info, cid: str) -> any:
    """
    查找对应的分P
    
    参数：
        video_info: 视频信息
        cid (str): 分P ID
        
    返回：
        any: 分P信息
    """
    logger = Logger()
    logger.info(f"开始查找分P: {cid}")
    logger.info(f"可用的分P列表: {[(p.cid, p.title) for p in video_info.parts]}")
    
    # 尝试不同的类型转换
    try:
        cid_int = int(cid)
        for p in video_info.parts:
            if int(p.cid) == cid_int:
                return p
    except ValueError:
        # 如果转换失败，直接使用字符串比较
        for p in video_info.parts:
            if str(p.cid) == str(cid):
                return p
    
    available_cids = [str(p.cid) for p in video_info.parts]
    error_msg = f"未找到分P: {cid}, 可用的分P: {available_cids}"
    logger.error(error_msg)
    raise ValueError(error_msg)

def find_quality(qualities: List, quality: str) -> any:
    """
    查找对应的清晰度
    
    参数：
        qualities (List): 清晰度列表
        quality (str): 目标清晰度
        
    返回：
        any: 清晰度信息
    """
    logger = Logger()
    logger.info(f"开始查找清晰度: {quality}")
    logger.info(f"可用的清晰度列表: {[(q.quality, q.height) for q in qualities]}")
    
    # 标准清晰度映射
    standard_qualities = {
        "1080P": 1080,
        "720P": 720,
        "480P": 480,
        "360P": 360
    }
    
    try:
        # 1. 直接匹配quality值
        for q in qualities:
            if str(q.quality) == str(quality):
                return q
        
        # 2. 如果没找到，尝试匹配标准清晰度
        target_height = standard_qualities.get(quality, None)
        if target_height:
            # 找到最接近的标准清晰度
            closest_quality = min(
                [q for q in qualities if q.height <= target_height],
                key=lambda x: abs(x.height - target_height),
                default=None
            )
            if closest_quality:
                logger.info(f"使用标准清晰度: {closest_quality.height}P")
                return closest_quality
        
        # 3. 如果还没找到，使用最接近的较低清晰度
        quality_height = int(quality)
        # 只考虑不高于目标清晰度的选项
        lower_qualities = [q for q in qualities if q.height <= quality_height]
        if lower_qualities:
            closest_quality = max(lower_qualities, key=lambda x: x.height)
            logger.info(f"使用最接近的较低清晰度: {closest_quality.height}P")
            return closest_quality
    
    except ValueError:
        # 如果转换失败，直接使用字符串比较
        for q in qualities:
            if str(q.quality) == str(quality):
                return q
    
    available_qualities = [f"{q.height}P (quality: {q.quality})" for q in qualities]
    error_msg = f"未找到清晰度: {quality}, 可用的清晰度: {available_qualities}"
    logger.error(error_msg)
    raise ValueError(error_msg)

async def download_with_retry(
    downloader: Downloader,
    url: str,
    output_path: str,
    task_id: str,
    part: str,
    max_retries: int = 3
) -> None:
    """
    带重试的下载
    
    参数：
        downloader (Downloader): 下载器
        url (str): 下载地址
        output_path (str): 输出路径
        task_id (str): 任务ID
        part (str): 下载部分（video/audio）
        max_retries (int): 最大重试次数
    """
    logger = Logger()
    retry_count = 0
    while retry_count < max_retries:
        try:
            if not url:
                raise ValueError(f"{part} URL为空")
            
            logger.info(f"{part}下载URL: {url}")
            await downloader.download_video(
                url,
                output_path,
                progress_callback=lambda d, t: update_progress(task_id, part, d, t)
            )
            break
        except Exception as e:
            retry_count += 1
            if retry_count >= max_retries:
                logger.error(f"下载{part}失败: {str(e)}")
                raise
            logger.warning(f"下载{part}失败，正在重试 ({retry_count}/{max_retries}): {str(e)}")
            await asyncio.sleep(1)

async def download_danmaku(
    downloader: Downloader,
    danmaku_url: str,
    task_id: str,
    max_retries: int = 3
) -> None:
    """
    下载弹幕
    
    参数：
        downloader (Downloader): 下载器
        danmaku_url (str): 弹幕URL
        task_id (str): 任务ID
        max_retries (int): 最大重试次数
    """
    logger = Logger()
    danmaku_path = DOWNLOAD_DIR / f"{task_id}_danmaku.xml"
    logger.info(f"开始下载弹幕: {danmaku_url}")
    TASKS[task_id]["current_step"] = "正在下载弹幕"
    
    retry_count = 0
    while retry_count < max_retries:
        try:
            await downloader.download_danmaku(danmaku_url, str(danmaku_path))
            break
        except Exception as e:
            retry_count += 1
            if retry_count >= max_retries:
                logger.warning(f"下载弹幕失败: {str(e)}")
                break
            logger.warning(f"下载弹幕失败，正在重试 ({retry_count}/{max_retries}): {str(e)}")
            await asyncio.sleep(1)
    
    if os.path.exists(danmaku_path):
        # 转换弹幕为ASS格式
        ass_path = DOWNLOAD_DIR / f"{task_id}_danmaku.ass"
        logger.info("开始转换弹幕为ASS格式")
        TASKS[task_id]["current_step"] = "正在转换弹幕格式"
        try:
            with open(danmaku_path, 'r', encoding='utf-8') as f:
                danmaku_list = DanmakuUtil.parse_xml(f.read())
            DanmakuUtil.convert_to_ass(danmaku_list, str(ass_path))
        except Exception as e:
            logger.warning(f"转换弹幕格式失败: {str(e)}")

async def download_subtitle(
    downloader: Downloader,
    subtitle: any,
    task_id: str,
    max_retries: int = 3
) -> None:
    """
    下载字幕
    
    参数：
        downloader (Downloader): 下载器
        subtitle: 字幕信息
        task_id (str): 任务ID
        max_retries (int): 最大重试次数
    """
    logger = Logger()
    sub_path = DOWNLOAD_DIR / f"{task_id}_sub_{subtitle.language}.json"
    logger.info(f"开始下载字幕: {subtitle.url}")
    TASKS[task_id]["current_step"] = f"正在下载{subtitle.language}字幕"
    
    retry_count = 0
    while retry_count < max_retries:
        try:
            await downloader.download_subtitle(subtitle.url, str(sub_path))
            break
        except Exception as e:
            retry_count += 1
            if retry_count >= max_retries:
                logger.warning(f"下载字幕失败: {str(e)}")
                break
            logger.warning(f"下载字幕失败，正在重试 ({retry_count}/{max_retries}): {str(e)}")
            await asyncio.sleep(1)
    
    if os.path.exists(sub_path):
        # 转换字幕为SRT格式
        srt_path = DOWNLOAD_DIR / f"{task_id}_sub_{subtitle.language}.srt"
        logger.info(f"开始转换字幕为SRT格式: {subtitle.language}")
        TASKS[task_id]["current_step"] = f"正在转换{subtitle.language}字幕格式"
        try:
            with open(sub_path, 'r', encoding='utf-8') as f:
                sub_list = SubtitleUtil.parse_json(f.read())
            SubtitleUtil.convert_to_srt(sub_list, str(srt_path))
        except Exception as e:
            logger.warning(f"转换字幕格式失败: {str(e)}")

def safe_filename(name: str) -> str:
    """
    生成安全的文件名
    
    参数：
        name (str): 原始文件名
        
    返回：
        str: 安全的文件名
    """
    return re.sub(r'[\\\\/:*?"<>|]', '_', name)

def get_unique_filename(base_name: str) -> Path:
    """
    获取唯一的文件名
    
    参数：
        base_name (str): 基础文件名
        
    返回：
        Path: 唯一的文件路径
    """
    final_name = safe_filename(base_name)
    final_path = DOWNLOAD_DIR / final_name

    # 如果目标文件已存在，避免覆盖
    if final_path.exists():
        base, ext = os.path.splitext(final_name)
        i = 1
        while (DOWNLOAD_DIR / f"{base}_{i}{ext}").exists():
            i += 1
        final_path = DOWNLOAD_DIR / f"{base}_{i}{ext}"
    
    return final_path

async def download_task(
    task_id: str,
    url: str,
    quality: str,
    cid: str,
    danmaku: bool,
    subtitle: bool
) -> None:
    """
    后台下载任务
    
    参数：
        task_id (str): 任务ID
        url (str): 视频URL
        quality (str): 视频质量
        cid (str): 分P ID
        danmaku (bool): 是否下载弹幕
        subtitle (bool): 是否下载字幕
    """
    logger = Logger()
    video_path = None
    audio_path = None
    
    try:
        # 初始化任务状态
        init_download_task(task_id)
        
        # 从URL中解析bvid
        bvid = await parse_bvid(url)
        logger.info(f"解析到视频ID: {bvid}")
        
        # 解析视频信息
        parser = Parser(url, Config())
        video_info = await parser.parse()
        
        # 查找对应的分P
        part = find_part(video_info, cid)
        logger.info(f"找到分P: {part.title} (CID: {part.cid})")
        
        # 获取播放信息
        video_url, audio_url, qualities = await parser.get_video_play_info(bvid, str(part.cid))
        logger.info(f"获取到视频播放地址: {video_url}")
        logger.info(f"获取到音频播放地址: {audio_url}")
        
        # 去重清晰度列表
        unique_qualities = []
        seen_heights = set()
        for q in qualities:
            if q.height not in seen_heights:
                seen_heights.add(q.height)
                unique_qualities.append(q)
        
        # 按清晰度高度排序
        qualities = sorted(unique_qualities, key=lambda x: x.height, reverse=True)
        logger.info(f"去重后的清晰度列表: {[(q.height, q.quality) for q in qualities]}")
        
        # 查找对应的清晰度
        quality_info = find_quality(qualities, quality)
        logger.info(f"找到清晰度: {quality_info.height}P (Quality: {quality_info.quality})")
        
        # 创建下载器
        downloader = Downloader()
        
        # 下载视频
        video_path = DOWNLOAD_DIR / f"{task_id}_video.mp4"
        logger.info(f"开始下载视频: {quality_info.url}")
        TASKS[task_id]["current_step"] = "正在下载视频"
        await download_with_retry(downloader, quality_info.url, str(video_path), task_id, "video")
        
        # 下载音频
        audio_path = DOWNLOAD_DIR / f"{task_id}_audio.m4a"
        logger.info(f"开始下载音频: {audio_url}")
        TASKS[task_id]["current_step"] = "正在下载音频"
        await download_with_retry(downloader, audio_url, str(audio_path), task_id, "audio")
        
        # 合并视频和音频
        output_path = DOWNLOAD_DIR / f"{task_id}.mp4"
        logger.info(f"开始合并视频和音频: {output_path}")
        TASKS[task_id]["current_step"] = "正在合并视频和音频"
        try:
            merger = Merger()
            merger.merge_video_audio(str(video_path), str(audio_path), str(output_path))
        except Exception as e:
            logger.error(f"合并视频和音频失败: {str(e)}")
            raise
        
        # 下载弹幕
        if danmaku and part.danmaku_url:
            await download_danmaku(downloader, part.danmaku_url, task_id)
        
        # 下载字幕
        if subtitle and part.subtitles:
            for sub in part.subtitles:
                await download_subtitle(downloader, sub, task_id)
        
        # 清理临时文件
        logger.info("清理临时文件")
        TASKS[task_id]["current_step"] = "正在清理临时文件"
        try:
            if os.path.exists(video_path):
                os.remove(video_path)
            if os.path.exists(audio_path):
                os.remove(audio_path)
        except Exception as e:
            logger.warning(f"清理临时文件失败: {str(e)}")
        
        # 更新任务状态
        logger.info("下载任务完成")
        TASKS[task_id]["status"] = "finished"
        TASKS[task_id]["progress"] = 100
        TASKS[task_id]["current_step"] = "下载完成"
        
        # 重命名文件
        final_path = get_unique_filename(f"{part.title}.mp4")
        os.rename(output_path, final_path)

        # 记录最终文件名，供前端下载接口使用
        TASKS[task_id]['final_filename'] = final_path.name
        
        # 插入文件结果
        insert_result(task_id, final_path.name, str(final_path), os.path.getsize(final_path), "video/mp4")
        
        # 更新数据库中的任务状态
        update_task_status(task_id, "finished", {
            "output_path": str(final_path),
            "file_name": final_path.name
        })
        
    except Exception as e:
        logger.error(f"下载任务失败: {str(e)}")
        TASKS[task_id]["status"] = "error"
        TASKS[task_id]["error"] = str(e)
        TASKS[task_id]["current_step"] = f"下载失败: {str(e)}"
        
        # 更新数据库中的任务状态
        update_task_status(task_id, "error", {"error": str(e)})
        
        # 清理临时文件
        try:
            if video_path and os.path.exists(video_path):
                os.remove(video_path)
            if audio_path and os.path.exists(audio_path):
                os.remove(audio_path)
        except:
            pass

async def parse_video_info(url: str) -> dict:
    """
    解析视频信息
    
    参数：
        url (str): 视频URL
        
    返回：
        dict: 包含视频信息的字典
    """
    try:
        logger = Logger()
        logger.info(f"开始解析视频: {url}")
        
        config = Config()
        parser = Parser(url, config)
        video_info = await parser.parse()
        
        logger.info(f"解析到视频信息: {video_info.title}")
        logger.info(f"分P数量: {len(video_info.parts)}")
        
        # 清晰度映射表
        quality_map = {
            1080: {"value": "232818", "name": "1080P 高清"},
            720: {"value": "171654", "name": "720P 准高清"},
            480: {"value": "103211", "name": "480P 标清"},
            360: {"value": "206230", "name": "360P 流畅"},
        }
        qualities = []
        seen_heights = set()
        for q in video_info.qualities:
            if q.height not in seen_heights:
                seen_heights.add(q.height)
                q_info = quality_map.get(q.height)
                if q_info:
                    qualities.append({
                        "value": q_info["value"],
                        "name": q_info["name"],
                        "url": q.url
                    })
        
        # 转换视频信息为前端需要的格式
        return {
            "title": video_info.title,
            "author": video_info.author,
            "cover": video_info.cover,
            "parts": [
                {
                    "cid": part.cid,
                    "title": f"P{i+1} {part.title}",
                    "video_url": part.video_url,
                    "audio_url": part.audio_url
                }
                for i, part in enumerate(video_info.parts)
            ],
            "qualities": qualities
        }
    except Exception as e:
        logger.error(f"解析视频失败: {str(e)}")
        raise 